# 导入StructuredTool 库  提供比@tool 更多的配置性
import asyncio

from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field


# 定义普通方法
def multiply(a: int, b: int) -> int:
    """ 两个数相乘 """
    return a * b


# 定义普通方法异步方法
async def async_multiply(a: int, b: int) -> int:
    """ 两个数相乘 """
    return a * b


# 方式1：简单配置
async def main():
    # func 参数： 指定一个同步函数，当你再同步上下文中调用工具时，它会使用这个同步函数来执行操作
    # corutine 参数： 指定一个异步函数。当你在异步上下文中调用工具时，它会使用这个异步函数来执行操作
    calculator = StructuredTool.from_function(func=multiply, coroutine=async_multiply)
    print(f"同步调用{calculator.invoke({"a": 2, "b": 3})}")
    print(f"异步调用{await calculator.ainvoke({"a": 2, "b": 6})}")


# 运行异步主函数 因为main里有异步调用 所以调用只能用异步调用
asyncio.run(main())


# 方式2：更多自定义参数配置
# 定义防具方法 方式1。2 ： 加@tool注解 定义字段描述 表示是langchain的工具方法  最终 方法名.args里 会多每个参数的描述
class CalculatorInput(BaseModel):
    a: int = Field(description="first number")
    b: int = Field(description="second number")


async def main2():
    calculator = StructuredTool.from_function(
        func=multiply,
        name="Calculator",
        description="两个数相乘",
        args_schema=CalculatorInput,
        return_type=True,
        # coroutine=async_multiply
    )
    print(calculator.invoke({"a": 2, "b": 33}))


asyncio.run(main2())
